#include "atomreference.h"

Atom::reference::reference() : atom(nullptr), offset(arma::fill::zeros), index(0) { }
Atom::reference::reference(Atom* const &a, const ivec4 &off, size_t i) : atom(a), offset(off), index(i) { }

long Atom::reference::i() const { return atom->i() + offset(0); }
long Atom::reference::j() const { return atom->j() + offset(1); }
long Atom::reference::h() const { return atom->h() + offset(2); }
long Atom::reference::l() const { return atom->l() + offset(3); }

ivec4 Atom::reference::ijhl() const { return atom->toIVec4() + offset; }
	
double Atom::reference::x() const { return Atom(atom->toIVec4() + offset).x(); }
double Atom::reference::y() const { return Atom(atom->toIVec4() + offset).y(); }
double Atom::reference::z() const { return Atom(atom->toIVec4() + offset).z(); }
	
vec3 Atom::reference::xyz() const { return Atom(atom->toIVec4() + offset).xyz(); }
vec2 Atom::reference::xy() const { return Atom(atom->toIVec4() + offset).xy(); }

// These define the grid-cell dimensions.
const double Atom::reference::grid::dx = 0.75;
const double Atom::reference::grid::dy = 0.75*sqrt(3.);

// The default constructor, use setAtoms() to set the list onto which this grid should be calculated.
Atom::reference::grid::grid() : atoms(nullptr), offsets(nullptr) { }
// Overload constructor, this set the list, and if the list is not empty generates the grid.
Atom::reference::grid::grid(Atom::list* const &list, imat* const &offs) : atoms(list), offsets(offs) { (*this)(); }
  
// This set the list, and if the list is not empty generates the grid.
void Atom::reference::grid::setAtoms(Atom::list* const &list, imat* const &offs) { atoms = list; offsets = offs; (*this)(); }

// If the Atom::list reference is set this re-generates the grid.
void Atom::reference::grid::operator()()
{
	if(!atoms || atoms->size() == 0)
		return;

	using namespace arma;

	const size_t origsize = atoms->size();
	const size_t perdim = (offsets) ? offsets->n_cols : 0;
	const size_t fullsize = origsize*(1+perdim);

	std::vector<arma::mat> XYs;
	XYs.resize(1+perdim,arma::mat(2,origsize));

	for(size_t a = 0; a < origsize; a++)
	{
		XYs[0].col(a) = atoms->at(a).xy();

		for(size_t p = 0; p < perdim; p++)
		{
			XYs[p+1].col(a) = Atom(atoms->at(a).toIVec4()+offsets->col(p)).xy();
		}
	}
	
	arma::vec mins = arma::min(XYs[0],1);
	arma::vec maxs = arma::max(XYs[0],1);
	
	double min_x = mins(0)-2*dx;
	double max_x = maxs(0)+2*dx;
	double min_y = mins(1)-2*dy;
	double max_y = maxs(1)+2*dy;
	
	origin(0) = min_x;
	origin(1) = min_y;
	
	long size_x = (max_x-min_x)/dx + 1;
	long size_y = (max_y-min_y)/dy + 1;
	
	grid2d.set_size(size_x,size_y);
	
	ivec4 zero = arma::zeros<ivec>(4);

	for(long a = 0; a < origsize; a++)
	{
		long i = (XYs[0](0,a)-min_x)/dx;
		long j = (XYs[0](1,a)-min_y)/dy;
		grid2d(i,j).push_back(reference(&((*atoms)[a]),zero,a));
		
		for(long p = 0; p < perdim; p++)
		{
			i = (XYs[p+1](0,a)-min_x)/dx;
			j = (XYs[p+1](1,a)-min_y)/dy;
			if( 0 <= i && i < size_x && 0 <= j && j < size_y)
				grid2d(i,j).push_back(reference(&((*atoms)[a]),offsets->col(p),a));
		}
	}
}
  
// This returns a subview of the cells surrounding the atom a.
const Atom::reference::list &Atom::reference::grid::getSurrounding(const Atom &a, long di, long dj)
{
	vec2 dxy = a.xy()-origin;
	long i = dxy(0)/dx + di;
	long j = dxy(1)/dy + dj;
	
	if(i < 0)
		i = 0;
	else if(i >= grid2d.n_rows)
		i = grid2d.n_rows-1;
	
	if(j < 0)
		j = 0;
	else if(j >= grid2d.n_cols)
		j = grid2d.n_cols-1;

	return grid2d(i,j);
}

void Atom::reference::grid::clear()
{
	grid2d.clear();
}